from genutils.multi_objective_ga import MogaPopulation
from genutils.io import from_config
from general_neuro.fast_thresh_detect import fast_thresh_detect
import random
import pylab
import numpy

def mystery_function(x):
    return trial_function(3.3, 2.4, 6.0, 1.5, x)

def trial_function(a, b, c, d, x):
    return a*x**numpy.sin(x**b) - c*numpy.cos(x) + x**d

def crossings(mys_fn, trial_fn, threshold = 20.0):
    mys_p, mys_n = fast_thresh_detect(mys_fn, threshold = threshold, refractory_period = 0)
    mys_crossings = len(mys_p) + len(mys_n)

    trial_p, trial_n = fast_thresh_detect(trial_fn, threshold = threshold, refractory_period = 0)
    trial_crossings = len(trial_p) + len(trial_n)

    return -abs(mys_crossings - trial_crossings)

def resid(mys_fn, trial_fn):
    return -numpy.sum((mys_fn - trial_fn)**2.0)

def mins(mys_fn, trial_fn):
    return -abs(numpy.min(mys_fn)- numpy.min(trial_fn))

def maxs(mys_fn, trial_fn):
    return -abs(numpy.max(mys_fn)- numpy.max(trial_fn))
    
class ThisPop(MogaPopulation):
    def score_model(self,gene,printme=False):
        a = gene[0]
        b = gene[1]
        c = gene[2]
        d = gene[3]
        # sample the mystery function at 150 pts from x=4 to 8
        mys_fn = numpy.array( [mystery_function(x) for x in numpy.linspace(4.0,8.0, 150)] )
        # sample the trial function at 150 pts from x=4 to 8
        trial_fn = numpy.array( [trial_function(a, b, c, d, x) for x in numpy.linspace(4.0,8.0, 150)] )
        scores = [crossings(mys_fn, trial_fn), crossings(mys_fn, trial_fn, threshold=30),
                  crossings(mys_fn, trial_fn, threshold=10),
                  resid(mys_fn, trial_fn),\
                  mins(mys_fn, trial_fn), maxs(mys_fn, trial_fn)]
        return scores

def plot_gene(gene):
    a = gene.allyls[0].value
    b = gene.allyls[1].value
    c = gene.allyls[2].value
    d = gene.allyls[3].value
    # sample the mystery function at 150 pts from x=4 to 8
    mys_fn = numpy.array( [mystery_function(x) for x in numpy.linspace(4.0,8.0, 150)] )
    # sample the trial function at 150 pts from x=4 to 8
    trial_fn = numpy.array( [trial_function(a, b, c, d, x) for x in numpy.linspace(4.0,8.0, 150)] )
    pylab.plot(mys_fn, color = 'black', linewidth=0.4)
    pylab.plot(trial_fn, linewidth=1.4) 
        
# read in parameters
cfile = "moga_test.parameters"
num_generations         = int(from_config('num_generations',cfile))
population_size         = int(from_config('population_size',cfile))
initial_population_size = int(from_config('initial_population_size',cfile))
mutate_rate             = from_config('mutate_rate',cfile)
random_mutate_rate      = from_config('random_mutate_rate',cfile)
cross_rate              = from_config('cross_rate',cfile)

# initalize the genetic algorithm population
thispop = ThisPop(population_size,initial_population_size,cfile)
thispop.mutate_rate = mutate_rate
thispop.cross_rate  = cross_rate

thispop.evolve(10, 20)
    

